# @Version        : 1.0
# @Update Time    : 2025/10/13 22:33
# @File           : schemes.py
# @IDE            : PyCharm
# @Desc           : 文件描述信息
import logging
import secrets
from typing import Optional, Dict, Sequence, Tuple, Callable
from urllib.parse import urlparse

from fastapi import HTTPException
from fastapi import Request
from fastapi.security import OAuth2PasswordBearer
from fastapi.security.utils import get_authorization_scheme_param
from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN

logger = logging.getLogger(__name__)


class MultiSourceOAuth2Bearer(OAuth2PasswordBearer):
    """多来源 Bearer 提取器（增强版）
    支持 header/cookie/query，提供 CSRF 校验（用于 cookie 源），限制 query 方法
    """

    def __init__(
        self,
        token_url: str,
        scheme_name: Optional[str] = None,
        scopes: Optional[Dict[str, str]] = None,
        description: Optional[str] = None,
        *,
        # 行为
        auto_error: bool = True,
        is_read_header: bool = True,
        is_read_cookie: bool = False,
        is_read_query: bool = False,
        # 各来源细节
        authorization: str = "Authorization",
        scheme: str = "Bearer",
        cookie_name: str = "access-token",
        query_name: str = "access_token",
        # 安全与限制
        require_csrf_for_cookies: bool = False,
        csrf_header_name: str = "X-CSRF-Token",
        csrf_cookie_name: str = "csrf-token",
        csrf_protected_methods: Sequence[str] = ("POST", "PUT", "PATCH", "DELETE"),
        allowed_query_methods: Sequence[str] = ("GET",),  # 默认仅 GET 允许用 query 携带
        csrf_validator: Optional[Callable[[Request], bool]] = None,  # 自定义校验钩子
        # 可选强化项（默认关闭以免误杀）
        check_origin: bool = False,
        allowed_origins: Optional[
            Sequence[str]
        ] = None,  # e.g. ["https://app.example.com"]
    ):
        self.is_read_header = is_read_header
        self.is_read_cookie = is_read_cookie
        self.is_read_query = is_read_query
        # 判断是否全部 False
        if not any((self.is_read_header, self.is_read_cookie, self.is_read_query)):
            self.is_read_header = True
        self.auto_error = auto_error

        self.authorization = authorization
        self.scheme = scheme
        self.cookie_name = cookie_name
        self.query_name = query_name

        # 安全配置
        self.require_csrf_for_cookies = require_csrf_for_cookies
        self.csrf_header_name = csrf_header_name
        self.csrf_cookie_name = csrf_cookie_name
        self.csrf_protected_methods = tuple(m.upper() for m in csrf_protected_methods)
        self.allowed_query_methods = tuple(m.upper() for m in allowed_query_methods)
        self.csrf_validator = csrf_validator
        self.check_origin = check_origin
        self.allowed_origins = tuple(allowed_origins) if allowed_origins else None

        super().__init__(
            tokenUrl=token_url,
            scheme_name=scheme_name,
            scopes=scopes,
            description=description,
            auto_error=False,  # 统一由本类控制抛错
        )

    async def __call__(self, request: Request) -> Optional[str]:
        try:
            token, source = self._get_token_and_source(request)
            if not token:
                if self.auto_error:
                    raise self._unauthorized("未提供有效的认证凭据")
                return None

            if (
                source == "cookie"
                and self.require_csrf_for_cookies
                and request.method.upper() in self.csrf_protected_methods
            ):
                ok = (
                    self.csrf_validator(request)
                    if self.csrf_validator is not None
                    else self._default_csrf_validate(request)
                )
                if not ok:
                    raise self._forbidden("CSRF 校验未通过")

            return token

        except HTTPException:
            raise
        except Exception as e:
            if self.auto_error:
                raise self._unauthorized(f"认证过程中发生错误: {str(e)}")
            return None

    def _get_token_and_source(
        self, request: Request
    ) -> Tuple[Optional[str], Optional[str]]:
        if self.is_read_header:
            token = self._from_header(request)
            src = "header"
        elif self.is_read_cookie:
            token = self._from_cookie(request)
            src = "cookie"
        else:
            token = self._from_query(request)
            src = "query"
        if token:
            return token, src
        return None, src

    def _from_header(self, request: Request) -> Optional[str]:
        authorization = request.headers.get(self.authorization)
        scheme, param = get_authorization_scheme_param(authorization)
        if scheme and param and scheme.lower() == self.scheme.lower():
            return param.strip() or None
        return None

    def _from_cookie(self, request: Request) -> Optional[str]:
        token = request.cookies.get(self.cookie_name)
        if token:
            token = token.strip()
            if token:
                return token
        return None

    def _from_query(self, request: Request) -> Optional[str]:
        if request.method.upper() not in self.allowed_query_methods:
            return None
        token = request.query_params.get(self.query_name)
        if token:
            token = token.strip()
            if token:
                return token
        return None

    def _unauthorized(self, detail: str) -> HTTPException:
        # 更规范的 WWW-Authenticate，可按需补充 error / error_description
        return HTTPException(
            status_code=HTTP_401_UNAUTHORIZED,
            detail=detail,
            headers={"WWW-Authenticate": self.scheme},
        )

    def _forbidden(self, detail: str) -> HTTPException:
        return HTTPException(
            status_code=HTTP_403_FORBIDDEN,
            detail=detail,
            headers={"WWW-Authenticate": self.scheme},
        )

    def _default_csrf_validate(self, request: Request) -> bool:
        # 1) Double Submit Cookie
        h = (request.headers.get(self.csrf_header_name) or "").strip()
        c = (request.cookies.get(self.csrf_cookie_name) or "").strip()
        if not h or not c:
            return False
        if not secrets.compare_digest(h, c):
            return False

        # 2) 可选：Origin/Referer 同源校验（按需开启）
        if self.check_origin:
            origin = (request.headers.get("Origin") or "").strip()
            if not origin:
                # 某些场景只带 Referer
                origin = (request.headers.get("Referer") or "").strip()
            if not origin:
                return False
            try:
                o = urlparse(origin)
                # 允许的来源列表或与当前服务同源
                if self.allowed_origins:
                    if origin.rstrip("/") not in [
                        x.rstrip("/") for x in self.allowed_origins
                    ]:
                        return False
                else:
                    # 与当前请求同源
                    cur = request.url
                    if (
                        o.scheme != cur.scheme
                        or o.hostname != cur.hostname
                        or (o.port or self._default_port(o.scheme))
                        != (cur.port or self._default_port(cur.scheme))
                    ):
                        return False
            except Exception as e:
                logger.warning(f"Origin 校验发生错误: {str(e)}")
                return False

        return True

    @staticmethod
    def _default_port(scheme: str) -> int:
        return 443 if scheme.lower() == "https" else 80
